#!/usr/bin/python
# -*- coding: UTF-8 -*-
import numpy as np
from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator  # 对二进制分类的评估器,它期望两个输入列:原始预测值和标签
from pyspark.ml.evaluation import MulticlassClassificationEvaluator  # 多类分类的评估器,它期望两个输入列:预测和标签

from sklearn.metrics import confusion_matrix

feature_list = ['Topography', 'Geological structure', 'Stratigraphic lithology', 'Altitude', 'Slope', 'Aspect',
                'Surface coverage', 'Vegetation cover', 'Soil moisture', 'Precipitation', 'Surface temperature',
                'Soil temperature']
test_size = 0.2
spark = SparkSession.builder.appName('random_forest').getOrCreate()
# 读取数据
df = spark.read.csv('train.csv', inferSchema=True, header=True)
# 将多个列合并成向量列的特征转换器,即将表中各列用一个类似list表示，输出预测列为单独一列。
df_assembler = VectorAssembler(inputCols=feature_list, outputCol="features")
df = df_assembler.transform(df)
model_df = df.select(['features', 'Tag'])  # 选择用于模型训练的数据
train_df, test_df = model_df.randomSplit([1 - test_size, test_size])  # 训练数据和测试数据划分

print('-----------使用随机森林进行数据训练----------------')
# numTrees设置随机数的数量为50,还有其他参数：maxDepth 树深;返回的模型类型为：RandomForestClassificationModel
rf_classifier = RandomForestClassifier(labelCol='Tag', numTrees=50).fit(train_df)
rf_predictions = rf_classifier.transform(test_df)

# print('{}{}'.format('评估每个属性的重要性:', rf_classifier.featureImportances))  # 评估每个功能的重要性,

# 下面对训练好的随机森林，完成重要性评估
importances = rf_classifier.featureImportances
print("重要性：", importances)
indices = np.argsort(importances)[::-1]
x_columns_indices = []
for f in range(len(feature_list)):
    # 对于最后需要逆序排序，我认为是做了类似决策树回溯的取值，从叶子收敛到根，根部重要程度高于叶子。
    print("%2d) %-*s %f" % (f + 1, 30, feature_list[indices[f]], importances[int(indices[f])]))
    x_columns_indices.append(feature_list[indices[f]])

# print(x_columns_indices)
# print(x_columns.shape[0])
# print(x_columns)
# print(np.arange(x_columns.shape[0]))

rf_predictions.select(['probability', 'Tag', 'prediction']).show(10, False)

print("------计算预测的准确率与混淆矩阵------")
rf_accuracy = MulticlassClassificationEvaluator(labelCol='Tag', metricName='accuracy').evaluate(rf_predictions)
print('MulticlassClassificationEvaluator 随机森林测试的准确性： {0:.0%}'.format(rf_accuracy))

rf_auc = BinaryClassificationEvaluator(labelCol='Tag').evaluate(rf_predictions)
print('BinaryClassificationEvaluator 随机森林测试的准确性： {0:.0%}'.format(rf_auc))

predict_Y = np.array(rf_predictions.select('Tag').collect())
y_test = np.array(test_df.select('Tag').collect())
cm = confusion_matrix(y_test, predict_Y)
print(cm)
